import random

import numpy as np
import networkx as nx
import pandas as pd
import matplotlib.pyplot as plt
import torch
from const import DEVICE, DEVICE_CPU


# 构建典型图GBN的函数，暂且先用这个图做测试。使用pandas.dataframe，也方便改造变成能直接输入图神经网络的数据格式
# 返回这个dataframe会方便之后操作
def GBN(device=DEVICE):
    edge = pd.DataFrame()
    edge['source'] = [0, 0, 1, 1, 1, 2, 3, 3, 4, 4, 4, 5, 5, 6, 7, 7, 9, 9, 10, 10, 11, 12, 12, 13, 14, 15]
    edge['target'] = [2, 8, 2, 3, 4, 4, 4, 9, 8, 9, 10, 6, 8, 7, 8, 10, 10, 12, 11, 12, 13, 14, 16, 14, 15, 16]
    edge['weight'] = np.ones(26)

    G = nx.from_pandas_edgelist(edge, source='source', target='target', edge_attr='weight')
    edgelist = [edge['source'].values.tolist(), edge['target'].tolist()]
    edgelist = torch.tensor(edgelist, dtype=torch.int).to(device)

    return G, edgelist


def GEANT2(device=DEVICE):
    edge = pd.DataFrame()
    edge['source'] = [0, 0, 1, 1, 1, 2, 2, 3, 3, 4, 5, 6, 6, 7, 7, 8, 8, 8, 8, 8, 9, 9, 9, 10, 11, 11, 12, 12, 14, 15,
                      16, 17, 18, 19, 21, 22]
    edge['target'] = [1, 2, 3, 6, 9, 3, 4, 5, 6, 7, 8, 8, 9, 8, 11, 11, 12, 17, 18, 20, 10, 12, 13, 13, 14, 20, 13, 21,
                      15, 16, 17, 18, 21, 23, 22, 23]
    edge['weight'] = np.ones(36)

    G = nx.from_pandas_edgelist(edge, source='source', target='target', edge_attr='weight')
    edgelist = [edge['source'].values.tolist(), edge['target'].tolist()]
    edgelist = torch.tensor(edgelist, dtype=torch.int).to(device)

    return G, edgelist


# 根据上下限随机分配线路容量
def generate_capacity_matrix(G, floor, ceil):
    node_num = G.nodes.__len__()
    capacity_matrix = np.zeros((node_num, node_num))
    for item in G.edges(data=False):
        val = random.randint(floor, ceil)
        capacity_matrix[item[0]][item[1]] = val
        capacity_matrix[item[1]][item[0]] = val
    return capacity_matrix


# v准确来说不是链路的传输速率，而是反映了链路两端的路由器转发能力，根据包大小不同，转发时延也不同
def generate_v_matrix(G, floor, ceil):
    node_num = G.nodes.__len__()
    v_matrix = np.zeros((node_num, node_num))
    for item in G.edges(data=False):
        v_matrix[item[0]][item[1]] = random.randint(floor, ceil)
        v_matrix[item[1]][item[0]] = random.randint(floor, ceil)
    return v_matrix


def adj(G):
    adj = nx.adjacency_matrix(G).todense()
    return adj


def draw_graph(G):
    pos = nx.spring_layout(G)
    nx.draw(G, pos, with_labels=True, font_weight='bold')
    edge_labels = nx.get_edge_attributes(G, 'weight')
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
    plt.show()


# 返回所有点到其他点的最短路径，是一个按点序号排序的列表，元素为(点编号，最短路径字典)，目标点编号是字典的key
def weighted_dijkstra(G):
    path_generator = nx.all_pairs_dijkstra_path(G)
    short_paths = []
    for item in path_generator:
        short_paths.append(item)
    return sorted(short_paths)
    # print(short_paths)


'''
weighted_adj have been modified
the output should be a 1*edge tensor
which is ought to fit the GCN input 
as well as to be in accordance with the edge_index
'''


def weighted_adj(G, device=DEVICE):
    edge_labels = nx.get_edge_attributes(G, 'weight')
    weight = torch.tensor(list(edge_labels.values()), dtype=torch.float32).to(device)

    return weight


def softmax(x):
    return np.exp(x) / np.sum(np.exp(x))


# 根据传入的权重更新图权重，这里的权重应当与边对应
def weight_update(G, weights, device=DEVICE):
    # 下面是更新权值
    edge_labels = nx.get_edge_attributes(G, 'weight')
    val_list = list(edge_labels.values())
    new_value = torch.tensor(val_list, dtype=torch.float32).to(device) + weights
    # print('weight_value before softmax')
    # print(new_value)
    # print('***********************************************')
    # new_value = softmax(new_value)
    min_val = torch.min(new_value.detach())

    if min_val < 0:
        new_value = new_value - min_val * 2
    new_edge_weight = (new_value / torch.sum(new_value))[0]

    new_edge_weight.to(device)

    # print(new_edge_weight)
    edges = G.edges
    i = 0
    for item in edges:
        G.add_edge(item[0], item[1], weight=new_edge_weight[i].item())
        i = i + 1

    return G
